import random
import operator
from typing import List, Dict, TypedDict, Annotated, Union, Literal
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph, END
load_dotenv()
import os


# --- 1. 定义更加完善的全局状态 ---
class GameState(TypedDict):
    # 基础配置
    players: List[str]  # 玩家名单 ["A", "B", "C", "D"...]
    player_words: Dict[str, str]  # 词汇分配
    spy_player: str  # 谁是卧底 (记录下来用于判决)

    # 游戏进程控制
    history: Annotated[List[str], operator.add]  # 聊天记录 (自动追加)
    current_turn_idx: int  # 当前轮到第几个玩家发言
    round_count: int  # 当前是第几轮辩论

    # 投票相关
    votes: Dict[str, str]  # {投票人: 被投人}
    game_result: str  # 最终结果


# 初始化模型 (用 GPT-4o 效果最好，3.5 可能会玩脱)
llm = ChatOpenAI(model=os.getenv("LONGMAO_MODEL"),
                 base_url=os.getenv("LONGMAO_BASE_URL"),
                 api_key=os.getenv("LONGMAO_API_KEY"))


# --- 2. 节点逻辑 (Node Functions) ---

def init_game_node(state: GameState):
    """上帝节点：初始化游戏，分配词汇"""
    print(f"\n=== 🎲 游戏初始化 (玩家人数: {len(state['players'])}) ===")

    # 词库：[平民词, 卧底词]
    word_pairs = [("元宵", "汤圆"), ("眼镜", "墨镜"), ("麦当劳", "肯德基")]
    common_word, spy_word = random.choice(word_pairs)

    players = state["players"]
    spy = random.choice(players)  # 随机选一个卧底

    allocation = {}
    for p in players:
        allocation[p] = spy_word if p == spy else common_word

    print(f"🔑 上帝视角：平民词[{common_word}]，卧底词[{spy_word}]，卧底是 -> {spy}")

    return {
        "player_words": allocation,
        "spy_player": spy,
        "current_turn_idx": 0,
        "round_count": 1,
        "history": [],
        "votes": {}
    }


def speaker_node(state: GameState):
    """通用的发言节点：轮到谁，谁就说话"""
    current_player = state["players"][state["current_turn_idx"]]
    my_word = state["player_words"][current_player]
    round_num = state["round_count"]

    print(f"🎤 轮到 {current_player} 发言 (第{round_num}轮)...")

    # 动态构建 Prompt
    prompt = f"""
    你是 {current_player}。正在玩《谁是卧底》。
    你的词是：【{my_word}】。
    现在是第 {round_num} 轮描述。

    请根据历史发言，用一句话描述你的词。
    要求：
    1. 不能直接说出【{my_word}】。
    2. 既要让队友听懂，又要迷惑卧底（或者如果你是卧底，要假装自己是好人）。
    3. 如果是第2轮以后，可以适当反驳别人的质疑。

    历史记录：
    {state['history']}
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    msg = f"[{current_player}]: {response.content}"
    print(f"   -> {msg}")

    # 更新状态：历史记录+1，轮次指针+1
    return {
        "history": [msg],
        "current_turn_idx": state["current_turn_idx"] + 1
    }


def voting_node(state: GameState):
    """投票节点：所有人一起投票 (简化为顺序投票)"""
    print(f"\n🗳️ === 第 {state['round_count']} 轮描述结束，开始投票 ===")
    players = state["players"]
    history_str = "\n".join(state["history"])
    new_votes = {}

    for p in players:
        my_word = state["player_words"][p]
        # 让每个 Agent 独立思考投给谁
        prompt = f"""
        你是 {p}，你的词是【{my_word}】。
        基于以下对话，谁最像卧底？

        对话记录：
        {history_str}

        请只输出一个名字（在 {players} 中选择），不要标点符号。
        """
        res = llm.invoke([HumanMessage(content=prompt)])
        target = res.content.strip()
        # 简单的清洗逻辑，防止 AI 输出 "我觉得是A" 这种废话
        for valid_name in players:
            if valid_name in target:
                target = valid_name
                break

        new_votes[p] = target
        print(f"   🗳️ {p} 投给了 -> {target}")

    return {"votes": new_votes}


def judge_node(state: GameState):
    """判决节点：计票并决定是【结束】还是【加赛】"""
    votes = state["votes"]
    spy = state["spy_player"]

    # 统计票数
    counts = {}
    for target in votes.values():
        counts[target] = counts.get(target, 0) + 1

    print(f"📊 得票统计: {counts}")

    # 排序找票数最高的
    sorted_votes = sorted(counts.items(), key=lambda x: x[1], reverse=True)
    top_target, top_count = sorted_votes[0]

    # 检查是否有平票 (Tie)
    # 如果第一名和第二名票数一样（且不只一个人得票），则是平局
    is_tie = False
    if len(sorted_votes) > 1 and sorted_votes[0][1] == sorted_votes[1][1]:
        is_tie = True

    if is_tie:
        print(f"⚠️ 平票！进入下一轮辩论！")
        return {
            "round_count": state["round_count"] + 1,
            "current_turn_idx": 0,  # 重置发言顺序
            "history": [f"--- 系统提示：上一轮投票平局，进入第 {state['round_count'] + 1} 轮辩论 ---"]
        }

    # 有人被淘汰
    print(f"❌ {top_target} 被淘汰！")
    if top_target == spy:
        result = "🎉 平民胜利！成功找出了卧底！"
    else:
        result = f"💀 卧底胜利！傻乎乎的平民淘汰了队友 {top_target}！"

    return {"game_result": result}


# --- 3. 路由逻辑 (The Router) ---

def turn_router(state: GameState) -> Literal["speak", "vote"]:
    """路由：决定是下一个人说话，还是去投票"""
    # 如果当前索引 < 玩家总数，说明这一轮还没说完
    if state["current_turn_idx"] < len(state["players"]):
        return "speak"
    # 否则，大家都说完了，去投票
    return "vote"


def game_router(state: GameState) -> Literal["continue_loop", "end"]:
    """路由：决定是游戏结束，还是因为平票继续循环"""
    if state.get("game_result"):
        return "end"
    return "continue_loop"  # 如果没有结果（平票），回到 Loop


# --- 4. 构建 LangGraph 图 ---

workflow = StateGraph(GameState)

# 添加节点
workflow.add_node("Init", init_game_node)
workflow.add_node("Speaker", speaker_node)  # 注意：只有一个 Speaker 节点被复用！
workflow.add_node("Voting", voting_node)
workflow.add_node("Judge", judge_node)

# 设置入口
workflow.set_entry_point("Init")

# 设置边 (核心逻辑)
workflow.add_edge("Init", "Speaker")

# 条件边 1：Speaker 说完后，是换下一个人，还是去投票？
workflow.add_conditional_edges(
    "Speaker",
    turn_router,
    {
        "speak": "Speaker",  # 循环调用自己！
        "vote": "Voting"
    }
)

workflow.add_edge("Voting", "Judge")

# 条件边 2：判决后，是结束，还是回炉重造？
workflow.add_conditional_edges(
    "Judge",
    game_router,
    {
        "continue_loop": "Speaker",  # 平票 -> 回到 Speaker 继续说
        "end": END
    }
)

# 编译
app = workflow.compile()

# --- 5. 运行 (动态设置玩家人数) ---
if __name__ == "__main__":
    # 这里你可以随便改人数，3个，4个，5个都能跑
    initial_state = {
        "players": ["Player_A", "Player_B", "Player_C", "Player_D"],
        "history": []
    }

    final_state = app.invoke(initial_state)

    print(f"\n🏆 最终结果：{final_state['game_result']}")